﻿using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using OpenCvSharp;
using System.Drawing;
using System.Drawing.Imaging;

namespace OnnxOCRApp
{
    public class PaddleOCR
    {
        private readonly InferenceSession _detectionSession;
        private readonly InferenceSession _recognitionSession;
        private readonly string[] _charList; // 字符列表，用于识别

        public PaddleOCR(string detModelPath, string recModelPath, string dictPath)
        {
            // 加载检测模型
            _detectionSession = new InferenceSession(detModelPath);

            // 加载识别模型
            _recognitionSession = new InferenceSession(recModelPath);

            // 加载字典
            _charList = File.ReadAllLines(dictPath);
        }

        public List<OCRResult> Process(string imagePath)
        {
            using (var bitmap = new Bitmap(imagePath))
            {
                return Process(bitmap);
            }
        }

        public List<OCRResult> Process(Bitmap image)
        {
            // 1. 文本检测
            var textBoxes = DetectTextBoxes(image);

            // 2. 文本识别
            var results = new List<OCRResult>();
            foreach (var box in textBoxes)
            {
                var croppedImage = CropImage(image, box);
                var text = RecognizeText(croppedImage);
                results.Add(new OCRResult { Text = text, Box = box });
            }

            return results;
        }

        private List<float[]> DetectTextBoxes(Bitmap image)
        {
            // 图像预处理
            var inputTensor = PreprocessForDetection(image);

            // 运行检测模型
            var inputs = new List<NamedOnnxValue> { NamedOnnxValue.CreateFromTensor("x", inputTensor) };
            using var results = _detectionSession.Run(inputs);

            // 后处理检测结果
            var outputTensor = results.First().AsTensor<float>();
            return PostprocessDetection(outputTensor, image.Width, image.Height);
        }

        private string RecognizeText(Bitmap image)
        {
            // 图像预处理
            var inputTensor = PreprocessForRecognition(image);

            // 运行识别模型
            var inputs = new List<NamedOnnxValue> { NamedOnnxValue.CreateFromTensor("x", inputTensor) };
            using var results = _recognitionSession.Run(inputs);

            // 后处理识别结果
            var outputTensor = results.First().AsTensor<float>();
            return PostprocessRecognition(outputTensor);
        }

        private DenseTensor<float> PreprocessForDetection(Bitmap image)
        {
            // Instead of using a fixed ratio, let's resize to exact dimensions expected by the model
            int targetHeight = 640;
            int targetWidth = 640;

            // Create a new bitmap with the target dimensions
            using var resizedImage = new Bitmap(targetWidth, targetHeight);
            using (var graphics = Graphics.FromImage(resizedImage))
            {
                // Use high quality scaling
                graphics.InterpolationMode = System.Drawing.Drawing2D.InterpolationMode.HighQualityBicubic;
                graphics.DrawImage(image, 0, 0, targetWidth, targetHeight);
            }

            // 转换为张量
            var tensor = new DenseTensor<float>(new[] { 1, 3, targetHeight, targetWidth });

            using (var ms = new MemoryStream())
            {
                resizedImage.Save(ms, ImageFormat.Bmp);
                ms.Position = 0;
                using var mat = Mat.FromStream(ms, ImreadModes.Color);

                // 归一化并转换为CHW格式
                mat.ConvertTo(mat, MatType.CV_32FC3, 1.0 / 255.0);

                var mean = new Scalar(0.485, 0.456, 0.406);
                var std = new Scalar(0.229, 0.224, 0.225);

                // 减均值除以标准差
                var channels = mat.Split();
                for (int i = 0; i < 3; i++)
                {
                    Cv2.Subtract(channels[i], mean[i], channels[i]);
                    Cv2.Divide(channels[i], std[i], channels[i]);
                }

                // 填充张量
                for (int c = 0; c < 3; c++)
                {
                    for (int h = 0; h < targetHeight; h++)
                    {
                        for (int w = 0; w < targetWidth; w++)
                        {
                            tensor[0, c, h, w] = channels[c].At<float>(h, w);
                        }
                    }
                }
            }

            return tensor;
        }

        private List<float[]> PostprocessDetection(Tensor<float> outputTensor, int originalWidth, int originalHeight)
        {
            // 实现检测后处理逻辑
            // 这里需要根据PaddleOCR的检测模型输出格式进行处理
            // 通常包括阈值过滤、多边形生成、NMS等步骤

            // 简化版实现，实际应用中需要更复杂的处理
            var boxes = new List<float[]>();

            // 假设输出是概率图，需要进行二值化和轮廓提取
            // 这里仅为示例，实际实现需要根据模型输出调整
            int height = outputTensor.Dimensions[2];
            int width = outputTensor.Dimensions[3];

            // 创建二值图像
            using var probMat = new Mat(height, width, MatType.CV_32FC1);
            for (int h = 0; h < height; h++)
            {
                for (int w = 0; w < width; w++)
                {
                    probMat.At<float>(h, w) = outputTensor[0, 0, h, w];
                }
            }

            // 二值化
            // 修改二值化阈值，可能需要调整以获得更好的结果
            using var binaryMat = new Mat();
            Cv2.Threshold(probMat, binaryMat, 0.2, 1, ThresholdTypes.Binary);

            // 转换为8位图像
            using var binaryMat8U = new Mat();
            binaryMat.ConvertTo(binaryMat8U, MatType.CV_8UC1, 255);

            // 查找轮廓
            Cv2.FindContours(binaryMat8U, out OpenCvSharp.Point[][] contours, out _, RetrievalModes.External, ContourApproximationModes.ApproxSimple);

            // 计算原始图像和处理图像的比例
            float scaleX = (float)originalWidth / width;
            float scaleY = (float)originalHeight / height;

            // 处理每个轮廓
            foreach (var contour in contours)
            {
                if (contour.Length < 4) continue;
            
                // 获取最小外接矩形
                var rect = Cv2.MinAreaRect(contour);
                var points = Cv2.BoxPoints(rect);

                // 缩放回原始图像尺寸
                var box = new float[8];
                for (int i = 0; i < 4; i++)
                {
                    box[i * 2] = points[i].X * scaleX;
                    box[i * 2 + 1] = points[i].Y * scaleY;
                }

                boxes.Add(box);
            }

            // 对检测到的框进行排序，从上到下，从左到右
            boxes = boxes.OrderBy(b => b[1]).ThenBy(b => b[0]).ToList();
            
            return boxes;
        }

        private Bitmap CropImage(Bitmap image, float[] box)
        {
            // 从原始图像中裁剪文本区域
            // 使用透视变换将四边形区域转换为矩形

            // 创建源点和目标点
            var srcPoints = new Point2f[4];
            for (int i = 0; i < 4; i++)
            {
                srcPoints[i] = new Point2f(box[i * 2], box[i * 2 + 1]);
            }

            // 计算宽度和高度
            float width = Math.Max(
                Distance(srcPoints[0], srcPoints[1]),
                Distance(srcPoints[2], srcPoints[3])
            );

            float height = Math.Max(
                Distance(srcPoints[0], srcPoints[3]),
                Distance(srcPoints[1], srcPoints[2])
            );

            // 创建目标点
            var dstPoints = new Point2f[4]
            {
                new Point2f(0, 0),
                new Point2f(width, 0),
                new Point2f(width, height),
                new Point2f(0, height)
            };

            // 使用OpenCV进行透视变换
            using (var ms = new MemoryStream())
            {
                image.Save(ms, ImageFormat.Bmp);
                ms.Position = 0;
                using var srcMat = Mat.FromStream(ms, ImreadModes.Color);

                // 计算透视变换矩阵
                var perspectiveMatrix = Cv2.GetPerspectiveTransform(srcPoints, dstPoints);

                // 应用透视变换
                using var dstMat = new Mat();
                Cv2.WarpPerspective(srcMat, dstMat, perspectiveMatrix, new OpenCvSharp.Size((int)width, (int)height));

                // 转换回Bitmap
                using var outStream = new MemoryStream();
                dstMat.WriteToStream(outStream);
                outStream.Position = 0;
                return new Bitmap(outStream);
            }
        }

        private float Distance(Point2f p1, Point2f p2)
        {
            return (float)Math.Sqrt(Math.Pow(p2.X - p1.X, 2) + Math.Pow(p2.Y - p1.Y, 2));
        }

        private DenseTensor<float> PreprocessForRecognition(Bitmap image)
        {
            // Update height to match the model's expected dimensions
            int height = 48; // Changed from 32 to 48 as required by the model
            int width = 320; // Fixed width or calculate based on aspect ratio

            using var resizedImage = new Bitmap(image, width, height);

            // 转换为张量
            var tensor = new DenseTensor<float>(new[] { 1, 3, height, width });

            using (var ms = new MemoryStream())
            {
                resizedImage.Save(ms, ImageFormat.Bmp);
                ms.Position = 0;
                using var mat = Mat.FromStream(ms, ImreadModes.Color);

                // 归一化并转换为CHW格式
                mat.ConvertTo(mat, MatType.CV_32FC3, 1.0 / 255.0);

                var mean = new Scalar(0.5, 0.5, 0.5);
                var std = new Scalar(0.5, 0.5, 0.5);

                // 减均值除以标准差
                var channels = mat.Split();
                for (int i = 0; i < 3; i++)
                {
                    Cv2.Subtract(channels[i], mean[i], channels[i]);
                    Cv2.Divide(channels[i], std[i], channels[i]);
                }

                // 填充张量
                for (int c = 0; c < 3; c++)
                {
                    for (int h = 0; h < height; h++)
                    {
                        for (int w = 0; w < width; w++)
                        {
                            tensor[0, c, h, w] = channels[c].At<float>(h, w);
                        }
                    }
                }
            }

            return tensor;
        }

        // 添加一个新方法用于保存调试图像
        public void SaveDebugImage(string originalImagePath, string outputPath, List<OCRResult> results)
        {
            using (var image = new Bitmap(originalImagePath))
            using (var graphics = Graphics.FromImage(image))
            {
                // 设置绘图属性
                var pen = new System.Drawing.Pen(Color.Red, 2);
                var font = new System.Drawing.Font("Arial", 12);
                var brush = new SolidBrush(Color.Red);

                // 绘制每个文本框和文本
                foreach (var result in results)
                {
                    // 绘制文本框
                    var points = new System.Drawing.Point[4];
                    for (int i = 0; i < 4; i++)
                    {
                        points[i] = new System.Drawing.Point(
                            (int)result.Box[i * 2],
                            (int)result.Box[i * 2 + 1]
                        );
                    }

                    graphics.DrawPolygon(pen, points);

                    // 绘制文本
                    graphics.DrawString(result.Text, font, brush, points[0]);
                }

                // 保存图像
                image.Save(outputPath);
            }
        }

        private string PostprocessRecognition(Tensor<float> outputTensor)
        {
            // 处理识别模型的输出
            // 通常是CTC解码过程

            // 获取每个时间步的预测
            int batchSize = outputTensor.Dimensions[0];
            int sequenceLength = outputTensor.Dimensions[1];
            int numClasses = outputTensor.Dimensions[2];

            // 对每个时间步找到最可能的字符
            var result = new List<int>();
            int lastCharIndex = -1; // 用于CTC合并

            // 添加置信度阈值过滤
            float confidenceThreshold = 0.5f; // 可以调整这个阈值
            
            for (int t = 0; t < sequenceLength; t++)
            {
                // 找到当前时间步的最大概率字符
                int maxIndex = 0;
                float maxProb = outputTensor[0, t, 0];

                for (int c = 1; c < numClasses; c++)
                {
                    if (outputTensor[0, t, c] > maxProb)
                    {
                        maxProb = outputTensor[0, t, c];
                        maxIndex = c;
                    }
                }

                // CTC解码：合并重复字符并移除空白字符
                if (maxIndex != lastCharIndex && maxIndex != 0) // 0通常是空白字符
                {
                    result.Add(maxIndex);
                }

                lastCharIndex = maxIndex;
            }

            // 将索引转换为字符
            var text = string.Join("", result.Select(i => i < _charList.Length ? _charList[i] : ""));
            
            // 移除可能的噪声字符
            text = CleanupText(text);
            
            return text;
        }

        // 添加一个新方法用于清理文本
        private string CleanupText(string text)
        {
            // 移除特殊字符和不可见字符
            var cleanText = new string(text.Where(c => !char.IsControl(c) && c != '�').ToArray());
            
            // 如果文本太短，可能是噪声
            if (cleanText.Length < 2)
                return string.Empty;
                
            return cleanText;
        }
    }

    public class OCRResult
    {
        public string Text { get; set; }
        public float[] Box { get; set; } // 8个值，表示4个点的坐标 [x1,y1,x2,y2,x3,y3,x4,y4]
    }
}
